from LMaug import *
from LMdataset import cate2idx, Skdict, kp2idx,idx2cate,idx2kp,kpInCates

class LMAugTrain(object):
    def __init__(self):
        self.augment = Compose([
            ResizeImg(size=(336, 336)),
            RandomHflip(),
            GenGauMask(r=1),
            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __call__(self, *args):
        return self.augment(*args)


class LMAugVal(object):
    def __init__(self):
        self.augment = Compose([
            ResizeImg(size=(336, 336)),
            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __call__(self, *args):
        return self.augment(*args)


def drop_shitdata(df):
    for cate in cate2idx:
        cate_idx = cate2idx[cate]
        kp1_idx = Skdict[cate_idx][0]
        kp1 = idx2kp[kp1_idx]
        kp2_idx = Skdict[cate_idx][1]
        kp2 = idx2kp[kp2_idx]

        drop_idx = df[(df['image_category'] == cate) & (df[kp1] == df[kp2])].index
        df.drop(drop_idx, axis=0, inplace=True)
        print '%s : drop shit data: %d' % (cate, len(drop_idx))
    return df